from z3 import *
import numpy as np
import draw_solution
import config


def abs(x):
    return If(x >= 0,x,-x)





def wtsn_smt_model(n_transactions,tasks,latency_bound):
    # processed data
    n_tasks = len(tasks)
    H = int(np.lcm.reduce([tasks[i]['T'] for i in range(n_tasks)]))
    print("The hyper-period is {}".format(H))
    nJobs = [int(H / tasks[i]['T']) for i in range(n_tasks)]

    # solver
    sol = Solver()

    # create variables
    s = [[Int("s_%s_%s" % (i+1, k+1)) for k in range(nJobs[i])]
        for i in range(n_tasks)]

    # FH point choice variables(TX)
    freq_tx = [[Int("f_tx_%s_%s"%(i+1,k+1)) for k in range(nJobs[i]) ]
            for i in range(n_tasks)]


    # formulate constraints
    # relative time window constraints 1
    r = [[Int("r_%s_%s" % (i+1, k+1)) for k in range(nJobs[i])]
        for i in range(n_tasks)]
    d = [[Int("d_%s_%s" % (i+1, k+1)) for k in range(nJobs[i])]
        for i in range(n_tasks)]
    for i in range(n_tasks):
        for k in range(nJobs[i]):
            sol.add(r[i][k] == tasks[i]['r'] + k * tasks[i]['T'])
            sol.add(d[i][k] == tasks[i]['d'] + k * tasks[i]['T'])
            sol.add(r[i][k] <= s[i][k],
                    s[i][k] <= d[i][k] - tasks[i]['p'])

    # jitter constraints 2 and 3
    jit = [Int("jit_%s" % (i+1)) for i in range(n_tasks)]
    for i in range(n_tasks):
        for k in range(nJobs[i]):
            for l in range(k+1, nJobs[i]):
                sol.add(jit[i] >= abs((s[i][k] - (k-1) * tasks[i]['T']) - (s[i][l] - (l-1) * tasks[i]['T'])))
        sol.add(jit[i] <= tasks[i]['jit'])

    # precedence constraints 8
    for i in range(n_tasks):
        for j in tasks[i]['pred']:
            for k in range(nJobs[i]):
                sol.add(s[j][k] + tasks[j]['p'] <= s[i][k])

    # end-to-end latency constraints 9 and 10
    lat = [Int("lat_%s" % (i+1)) for i in range(n_transactions)]
    for i in range(n_transactions):
        sources = [t for t in range(n_tasks) if tasks[t]['transaction'] == i and not tasks[t]['pred']]
        destinations = [t for t in range(n_tasks) if tasks[t]['transaction'] == i and not tasks[t]['succ']]
        for source in sources:
            for dest in destinations:
                for k in range(nJobs[source]):
                    sol.add(lat[i] >= s[dest][k] + tasks[dest]['p'] - s[source][k])
        sol.add(lat[i] <= latency_bound[i])


    # resource constraints 11
    d = [tasks[i]['d'] + k * tasks[i]['T'] for k in range(nJobs[i]) for i in range(n_tasks)]
    n_H_max = int(np.ceil(max(d)/H))
    for i in range(n_tasks):
        for j in range(i, n_tasks):
            for k in range(nJobs[i]):
                for l in range(nJobs[j]):
                    for f in range(n_H_max):
                        for h in range(n_H_max):
                            if tasks[i]['fix'] == tasks[j]['fix']:
                                if i != j or k != l:
                                    sol.add(
                                        Xor(
                                            ((s[i][k] + f*H) + tasks[i]['p']) <= (s[j][l] + h*H),
                                            ((s[j][l] + h*H) + tasks[j]['p']) <= (s[i][k] + f*H)
                                        )
                                    )
                                               
    # freq contraints
    for i in range(n_tasks):
        for k in range(nJobs[i]):
            sol.add(freq_tx[i][k]>=0,freq_tx[i][k]<config.FREQ_NUM)
    
    # freq contraints, One freq_point only can be uesd by one task on the same time
    for f in range(config.FREQ_NUM):
        for i in range(n_tasks):
            for j in range(i+1,n_tasks):
                for k in range(nJobs[i]):
                    for l in range(nJobs[j]):
                        for m in range(n_H_max):
                            for h in range(n_H_max):
                                sol.add(Implies(
                                    And(
                                        freq_tx[i][k] == f,
                                        freq_tx[j][l] == f    
                                    ),
                                    Or(
                                        (s[i][k]+m*H+tasks[i]['p']) <= (s[j][l]+h*H),
                                        (s[j][l]+h*H+tasks[j]['p']) <= (s[i][k]+m*H)
                                    )
                                ))
    
    # freq contraints: 
    for i in range(n_tasks):
        for j in range(i+1,n_tasks):
            for k in range(nJobs[i]):
                for l in range(nJobs[j]):
                    for m in range(n_H_max):
                        for h in range(n_H_max):
                            # 1.one machine cannot be simultaneously use one FH to both tx and rx
                            if tasks[i]['fix'] == tasks[j]['next'][0] or tasks[j]['fix']== tasks[i]['next'][0]:
                                sol.add(Implies(
                                    s[i][k]+m*H == s[j][l]+h*H,
                                    freq_tx[i][k] != freq_tx[j][l]
                                ))
                            # 2. one machine cannot be simultaneously received two FHs 
                            # (while sending simultaneously is constrainted by constraints 11 )
                            if tasks[i]['next'][0] == tasks[j]['next'][0]:
                                sol.add(s[i][k]+m*H != s[j][l]+h*H)
                                
                            

    # freq constraints: 3.each machine use FIFO,i.e, if task_i process first, task_i's succ task must process first too
    for i in range(n_tasks):
        if not tasks[i]['succ']:
            continue
        next_task_i = tasks[i]['succ'][0]
        for j in range(i+1,n_tasks):
            if not tasks[j]['succ']:
                continue
            next_task_j = tasks[j]['succ'][0]
            if tasks[i]['next'][0] == tasks[j]['next'][0]:
                for k in range(nJobs[i]):
                    for l in range(nJobs[j]):
                        for m in range(n_H_max):
                            for h in range(n_H_max):
                                sol.add(Implies(
                                    s[i][k]+m*H > s[j][l]+h*H,
                                    s[next_task_i][k]+m*H > s[next_task_j][l]+h*H
                                ))
                                sol.add(Implies(
                                    s[i][k]+m*H < s[j][l]+h*H,
                                    s[next_task_i][k]+m*H < s[next_task_j][l]+h*H
                                ))
    
                    

    # solve the model
    if sol.check() == sat:
        print("Succeeded to solve")
        
        m = sol.model()
        start_times = [[m.evaluate(s[i][j]).as_long() for j in range(nJobs[i])]
                    for i in range(n_tasks)]
        print("\nMatrix of the start times is ")
        print_matrix(start_times)
        print("\nVector of task mappings to resources")
        print([tasks[i]['fix'] for i in range(n_tasks)])
        # print freq points
        freqs = [[m.evaluate(freq_tx[i][j]).as_long() for j in range(nJobs[i])]
                for i in range(n_tasks)]
        print("\n freq points mapping:")
        print_matrix(freqs)
        return start_times,freqs
    else:
        print("Failed to solve")
        return None
    

    return


if __name__ == "__main__":
    # input data -----------------------------------------------------
    n_transactions = 2
    tasks = {
        0: {'r': 0, 'd': 10, 'T': 2, 'jit':1, 'p':1, 'fix':1, 'pred': [], 'succ': [2], 'transaction':0, 'next':[2]},
        1: {'r': 0, 'd': 10, 'T': 15, 'jit':5, 'p':1, 'fix':1, 'pred': [], 'succ': [4], 'transaction':1, 'next':[2]},
        2: {'r': 0, 'd': 10, 'T': 2, 'jit':1, 'p':1, 'fix':2, 'pred': [0], 'succ': [3], 'transaction':0, 'next':[3]},
        3: {'r': 0, 'd': 10, 'T': 2, 'jit':1, 'p':1, 'fix':3, 'pred': [2], 'succ': [], 'transaction':0, 'next':[]},
        4: {'r': 0 , 'd': 10, 'T': 15, 'jit':5, 'p':1, 'fix':2, 'pred': [1], 'succ': [5], 'transaction':1, 'next':[3]},
        5: {'r': 0 , 'd': 10, 'T': 15, 'jit':5, 'p':1, 'fix':3, 'pred': [4], 'succ': [], 'transaction':1, 'next':[]}
        }
    latency_bound = [5, 10]
    
    start_times,freqs = wtsn_smt_model(n_transactions,tasks,latency_bound)
    
    # if start_times != None and freqs != None:
    #     # draw the sheduling result of WTSN
    #     draw_solution.draw_machines_res(tasks,start_times,freqs)
    
    




